package org.nicta.learn.factorization.GreedyResidualFittingFeatures;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Random;

import org.nicta.colt.utils.ColtUtils;
import org.nicta.datasetLoader.MovieLens.MovieLensDataLoader;
import org.nicta.learn.optimization.LBFGS;

import cern.colt.matrix.DoubleFactory1D;
import cern.colt.matrix.DoubleFactory2D;
import cern.colt.matrix.DoubleMatrix1D;
import cern.colt.matrix.DoubleMatrix2D;
import cern.colt.matrix.linalg.Algebra;
import cern.jet.math.Functions;

public class FeaturesGRF {
	
	MovieLensDataLoader mlLoader;
	private HashMap<Integer, double[]> userFeatures;
	//userFeaturesMat holds users features excluding the features encoding users
	//Ids in the colt 1D matrix format
	private HashMap<Integer, DoubleMatrix1D> userFeaturesMat;
	//userFullFeaturesMat holds complete users features 
	private HashMap<Integer, DoubleMatrix1D> userFullFeaturesMat;
	private HashMap<Integer, DoubleMatrix1D> userTraits;

	private HashMap<Integer, double[]> movieFeatures;
	//movieFeaturesMat holds movie features excluding the features encoding 
	//movies Id in the colt 1D matrix format
	private HashMap<Integer, DoubleMatrix1D> movieFeaturesMat;
	//movieFullFeaturesMat holds complete movies features 
	private HashMap<Integer, DoubleMatrix1D> movieFullFeaturesMat;
	private HashMap<Integer, DoubleMatrix1D> movieTraits;

	//ratings holds data for training
	private HashMap<Integer[], Double> ratings;
	//testRatings holds data for training
	private HashMap<Integer[], Double> testRatings;

	private HashMap<Integer, HashSet<Integer>> userRatedMovies;
	private HashMap<Integer, HashSet<Integer>> moviesUsersMap;
	private HashMap<Integer, HashMap<Integer, Double>> movieUserRatings;
	//private HashMap<String, DoubleMatrix2D> objectPool;
	
	
	private DoubleMatrix2D userFeatureProjection;
	private DoubleMatrix2D movieFeatureProjection;
	
	private DoubleMatrix2D userFeatureProjectionTemp;
	private DoubleMatrix2D movieFeatureProjectionTemp;
	
	private Algebra algebra;
	private double globalRatingMean = 0;
	private HashMap<Integer, Double> userAverage;
	private HashMap<Integer, Double> itemAverage;

	
	private Integer dimension; // latent dimension
	private double alpha; //Learning Rate
	private double lambdaUser, lambdaMovie; //regularization control
	private double momentum;
	private double testSetFraction;
	private double minStepSize = 1e-2;
	
	
	public FeaturesGRF(String datasetType, String ratingsFileLocation,
					  String usersFileLocation,String moviesFileLocation,
					  Integer dimension )throws Exception {
		
		this.dimension = dimension;
		
		mlLoader = new MovieLensDataLoader(datasetType, ratingsFileLocation,
				usersFileLocation, moviesFileLocation);
		//loadData
		mlLoader.loadData();
		userFeatures = mlLoader.getUsersFeatures();
		movieFeatures = mlLoader.getMoviesFeatures();
		ratings = mlLoader.getRatings();
		userRatedMovies = mlLoader.getUserRatedMoviesMap();
		moviesUsersMap = mlLoader.getMoviesUsersMap();
		movieUserRatings = mlLoader.getmovieUserRatings();
		
		initialize();
	}
	
	private void initialize()
	{
		 algebra = new Algebra();
		 userFeatureProjection = null;
		 movieFeatureProjection = null;
		 alpha = 0.01;
		 lambdaUser = 10;
		 lambdaMovie = 10;
		 momentum = 0.0;
		 testSetFraction = 0.1;
	
		 testRatings = createTestSet(ratings, moviesUsersMap, userRatedMovies,
				 testSetFraction);
		 
		 userAverage = getUserAverage(movieUserRatings,userRatedMovies);
		 ratingNormalizationPerUserAverage(ratings,userAverage);
		 
		 movieFeaturesMat = new HashMap<Integer, DoubleMatrix1D>();
		 movieFullFeaturesMat = new HashMap<Integer, DoubleMatrix1D>();
		 
		 userFeaturesMat = new HashMap<Integer, DoubleMatrix1D>();
		 userFullFeaturesMat = new HashMap<Integer, DoubleMatrix1D>();	 
		 
		 constructSparseMovieFeatureMatrices(movieFeatures,movieFeaturesMat,
				 movieFullFeaturesMat);
		 constructSparseUserFeatureMatrices(userFeatures, userFeaturesMat,
				 userFullFeaturesMat);
		 

	}
	
	private double calculateRatingMean(HashMap<Integer[], Double> ratings){
		double sum = 0;
		Object[] ratingsKey = ratings.keySet().toArray();
		for(Object userItemPair: ratingsKey){
			Double rating =  ratings.get(userItemPair);
			sum += rating;
		}
		return sum/ratingsKey.length;
	}
	
	private void ratingNormalizationPerUserAverage(
			HashMap<Integer[], Double> ratings,
			HashMap<Integer, Double> userAverage 
			){
		Object[] ratingsKey = ratings.keySet().toArray();
		for(Object userItemPair: ratingsKey){
			Double rating =  ratings.get(userItemPair);
			Integer userId = ((Integer[]) userItemPair)[0];
			Integer movieId = ((Integer[]) userItemPair)[1];
			rating = rating-userAverage.get(userId);
			ratings.put((Integer[]) userItemPair,rating);
			movieUserRatings.get(movieId).put(userId, rating);
		}
	}
	
	private HashMap<Integer, Double> getUserAverage(
			HashMap<Integer, HashMap<Integer,Double>> movieUserRatings,
			HashMap<Integer, HashSet<Integer>> userRatedMovies){
		double average;
		double sum;
		HashMap<Integer, Double> userAverage = new HashMap<Integer, Double>();
		for(Integer u: userRatedMovies.keySet()){
			HashSet<Integer> uRatedMovies = userRatedMovies.get(u);
			sum = 0;
			for(Integer m: uRatedMovies){
				sum += movieUserRatings.get(m).get(u);
			}
			average = sum/uRatedMovies.size();
			userAverage.put(u, average);
		}
		return userAverage;
	}
	
	
	private HashMap<Integer, Double> getItemAverage(
			HashMap<Integer, HashMap<Integer,Double>> movieUserRatings,
			HashMap<Integer, HashSet<Integer>> moviesUsersMap){
		double average;
		double sum;
		HashMap<Integer, Double> itemAverage = new HashMap<Integer, Double>();
		for(Integer m: moviesUsersMap.keySet()){
			HashSet<Integer> movieRatedUsers = moviesUsersMap.get(m);
			sum = 0;
			for(Integer u: movieRatedUsers){
				sum += movieUserRatings.get(m).get(u);
			}
			if(sum == 0) continue;//movies that are not rated
			average = sum/movieRatedUsers.size();
			itemAverage.put(m, average );
		}
		return itemAverage;
	}
	

	private void meanNormalization(HashMap<Integer[], Double> ratings){
		globalRatingMean = calculateRatingMean(ratings);
		for(Integer[] userItemPair: ratings.keySet()){
			Double rating =  ratings.get(userItemPair);
			Integer userId = userItemPair[0];
			Integer movieId = userItemPair[1];
			ratings.put(userItemPair,rating-globalRatingMean);
		}
		
	}
	
	private double calculatePredictedRating(Integer userId,Integer movieId){
		DoubleMatrix1D userFeature = 
				DoubleFactory1D.sparse.make(userFeatures.get(userId));
		DoubleMatrix1D movieFeature = 
				DoubleFactory1D.sparse.make(movieFeatures.get(movieId));
		DoubleMatrix1D Ux = algebra.mult(userFeatureProjection, userFeature);
		DoubleMatrix1D Vy = algebra.mult(movieFeatureProjection, movieFeature);
		return algebra.mult(Ux, Vy);
	}
		
	private HashMap<Integer[], Double> createTestSet(
			HashMap<Integer[], Double> ratings,
			HashMap<Integer, HashSet<Integer>>  moviesUsersMap,
			HashMap<Integer, HashSet<Integer>> userRatedMovies,
			double fraction){
		
		HashMap<Integer[], Double> testRatings;
		testRatings = new HashMap<Integer[], Double>();
		Object[] ratingsKey = ratings.keySet().toArray();
		double testSetSize = fraction * ratings.keySet().size();
		Random rand = new Random();
		int key;
		while(testRatings.size() < testSetSize){
			key = rand.nextInt(ratingsKey.length);
			Object userItemPair = ratingsKey[key];
			if(!ratings.containsKey(userItemPair)) continue;
			Integer userId = ((Integer[]) userItemPair)[0];
			Integer movieId = ((Integer[]) userItemPair)[1];
			if(moviesUsersMap.get(movieId).size() > 1 && 
					userRatedMovies.get(userId).size() > 1){

				Double rating =  ratings.get(userItemPair);
				ratings.remove(userItemPair);
				userRatedMovies.get(userId).remove(movieId);
				moviesUsersMap.get(movieId).remove(userId);
				testRatings.put(new Integer[]{userId, movieId}, rating);
			
			}
		}
		return testRatings;
	}
	
	private double[] calculateTestError(
			HashMap<Integer, DoubleMatrix1D> userTraits,
			HashMap<Integer, DoubleMatrix1D> movieTraits){
	
		double rmseCost = 0.0;
		double maeCost = 0.0;
		double predictedRating;
		Object[] ratingsKey  = testRatings.keySet().toArray();
		double [] result = new double[2];
		
		for(Object userItemPair: ratingsKey){
			Integer userId = ((Integer[]) userItemPair)[0];
			Integer movieId = ((Integer[]) userItemPair)[1];
			Double rating =  testRatings.get(userItemPair);
			predictedRating = algebra.mult(userTraits.get(userId), 
					movieTraits.get(movieId)) + userAverage.get(userId);
			
			//cutoff
			if (predictedRating > 5) predictedRating = 5;
			if (predictedRating < 1) predictedRating = 1;
			
			rmseCost += Math.pow((rating - predictedRating),2);
			maeCost += Math.abs(rating - predictedRating);
		}
		result[0] = Math.sqrt(rmseCost*(1.0/testRatings.size()));
		result[1] = maeCost*(1.0/testRatings.size());
		return result;
	}	

	

	private void constructSparseUserFeatureMatrices(
			HashMap<Integer, double[]> userFeatures, 
			HashMap<Integer, DoubleMatrix1D> userFeaturesMat,
			HashMap<Integer, DoubleMatrix1D> userFullFeaturesMat){
		//done only for  optimization

		for(Integer i : userFeatures.keySet()){
			DoubleMatrix1D x = DoubleFactory1D.sparse.make(userFeatures.get(i));
			DoubleMatrix1D y =  DoubleFactory1D.sparse.make(mlLoader.getUsersCount());
			y.set(i-1,1);
			userFullFeaturesMat.put(i,DoubleFactory1D.sparse.append(x, y));
			userFeaturesMat.put(i,x);
		}
	}
	
	private void constructSparseMovieFeatureMatrices(
			HashMap<Integer, double[]> movieFeatures, 
			HashMap<Integer, DoubleMatrix1D> movieFeaturesMat,
			HashMap<Integer, DoubleMatrix1D> movieFullFeaturesMat){
		//done only for speed optimization

		for(Integer i : movieFeatures.keySet()){
			DoubleMatrix1D x = DoubleFactory1D.sparse.make(movieFeatures.get(i));
			DoubleMatrix1D y =  DoubleFactory1D.sparse.make(mlLoader.getMoviesCount());
			y.set(i-1,1);
			movieFullFeaturesMat.put(i,DoubleFactory1D.sparse.append(x, y));
			movieFeaturesMat.put(i,x);
		}
	}
	
	/*
	 * Code Block for Calculating Gradients by Finite Difference Approximation 
	 * for verifying the correctness of the gradient (Sledge hammer for checking
	 * correctness of the gradient but freaking slow for large features) 
	 */

	private double calculateCost(HashMap<Integer[], Double> residue,
			DoubleMatrix2D userFeatureProjection,
			DoubleMatrix2D movieFeatureProjection,
			HashMap<Integer, DoubleMatrix1D> userTraits,
			HashMap<Integer, DoubleMatrix1D> movieTraits){
		
		double cost = 0.0;
		double norm = 0.0;
		double predictedRating;
		Object[] ratingsKey  = residue.keySet().toArray();
		for(Object userItemPair: ratingsKey){
			Integer userId = ((Integer[]) userItemPair)[0];
			Integer movieId = ((Integer[]) userItemPair)[1];
			Double rating =  residue.get(userItemPair);
			predictedRating = algebra.mult(userTraits.get(userId), 
					movieTraits.get(movieId));
			cost += Math.pow((rating - predictedRating),2);
		}
		norm = (Math.pow(algebra.normF(userFeatureProjection),2) * lambdaUser+ 
				Math.pow(algebra.normF(movieFeatureProjection),2)* lambdaMovie)/2;
		

//		cost = cost*(1.0/ratings.size()) + norm;
		cost = cost/2.0+norm;

		return cost;
	}

	
	
	private HashMap<Integer, DoubleMatrix1D> getUserTraits(
			DoubleMatrix2D userFeatureProjection){
		
		HashMap<Integer, DoubleMatrix1D> userTraits = new 
				HashMap<Integer, DoubleMatrix1D>();

		for(Integer userId: userFullFeaturesMat.keySet()){
			DoubleMatrix1D userTrait = algebra.mult(userFeatureProjection, 
					userFullFeaturesMat.get(userId));
			userTraits.put(userId, userTrait);
		}

		return userTraits;
	}
	
	private HashMap<Integer, DoubleMatrix1D> getMovieTraits(
			DoubleMatrix2D movieFeatureProjection){
		
		HashMap<Integer, DoubleMatrix1D> movieTraits = new 
				HashMap<Integer, DoubleMatrix1D>();
		for(Integer movieId: movieFullFeaturesMat.keySet()){
			DoubleMatrix1D movieTrait = algebra.mult(movieFeatureProjection, 
					movieFullFeaturesMat.get(movieId));
			movieTraits.put(movieId, movieTrait);
		}

		return movieTraits;
	}
	
	private DoubleMatrix2D getUserGradient(HashMap<Integer[], Double> residue, 
			DoubleMatrix2D userFeatureProjection,
			HashMap<Integer, DoubleMatrix1D>userTraits,
			HashMap<Integer, DoubleMatrix1D> movieTraits){
		
		DoubleMatrix2D userGradient = DoubleFactory2D.dense.make(1,
				 mlLoader.getUserFeaturesSize());
		
		DoubleMatrix2D userGradientUserId = DoubleFactory2D.dense.make(
				1,mlLoader.getUsersCount());
		
		DoubleMatrix1D userFeature;
		DoubleMatrix1D  Vy;
		DoubleMatrix1D  Ux;
		DoubleMatrix2D pointGradient= DoubleFactory2D.dense.make(1,
				mlLoader.getUserFeaturesSize());
		
		double predictedRating;
		double error;
		Object[] ratingsKey  = residue.keySet().toArray();
		for(Object userItemPair: ratingsKey){
				Integer userId = ((Integer[]) userItemPair)[0];
				Integer movieId = ((Integer[]) userItemPair)[1];
				Double rating =  residue.get(userItemPair);
				
				Ux = userTraits.get(userId);
				Vy = movieTraits.get(movieId);
				
				predictedRating = algebra.mult(Ux,Vy);
				error = rating - predictedRating ;
				userFeature = userFeaturesMat.get(userId);
				algebra.multOuter(Vy, userFeature, pointGradient);
				
				pointGradient.assign(Functions.mult(error));
				userGradient.assign(pointGradient,Functions.minus);
				
				ColtUtils.sumColumnByVector(userGradientUserId, 
						Vy.copy().assign(Functions.mult(-error)), userId-1);

		}
		
		userGradient = DoubleFactory2D.dense.appendColumns(userGradient, userGradientUserId);
		
		//add regulariation parameters
		userGradient.assign(ColtUtils.multScalar(
				userFeatureProjection, lambdaUser),Functions.plus);
		return userGradient;
	}
	
	private DoubleMatrix2D getMovieGradient(HashMap<Integer[], Double> residue,
			DoubleMatrix2D movieFeatureProjection,
			HashMap<Integer, DoubleMatrix1D>userTraits,
			HashMap<Integer, DoubleMatrix1D> movieTraits){
		
		
		DoubleMatrix2D movieGradient = DoubleFactory2D.dense.make(1,
				 mlLoader.getMovieFeaturesSize());
		
		DoubleMatrix2D movieGradientUserId = DoubleFactory2D.dense.make(
				1,mlLoader.getMoviesCount());
				
		DoubleMatrix1D movieFeature;
		DoubleMatrix1D  Ux;
		DoubleMatrix1D Vy;
		DoubleMatrix2D pointGradient= DoubleFactory2D.dense.make(1,
				mlLoader.getMovieFeaturesSize());
		
		double predictedRating;
		double error;
		Object[] ratingsKey  = residue.keySet().toArray();
		for(Object userItemPair: ratingsKey){
			Integer userId = ((Integer[]) userItemPair)[0];
			Integer movieId = ((Integer[]) userItemPair)[1];
			Double rating =  residue.get(userItemPair);
			
			Ux = userTraits.get(userId);
			Vy = movieTraits.get(movieId);
			
			predictedRating = algebra.mult(Ux,Vy);
			error = rating - predictedRating ;
			
			movieFeature = movieFeaturesMat.get(movieId);
			algebra.multOuter(Ux, movieFeature, pointGradient);
			pointGradient.assign(Functions.mult(error));
			movieGradient.assign(pointGradient,Functions.minus);
			
			//exploit the sparsity for speedup
			ColtUtils.sumColumnByVector(movieGradientUserId, 
					Ux.copy().assign(Functions.mult(-error)), movieId-1);
			
		}
				
		movieGradient = DoubleFactory2D.dense.appendColumns(movieGradient, movieGradientUserId);
		
		//add regularization parameters
		movieGradient.assign(ColtUtils.multScalar(
				movieFeatureProjection, lambdaMovie),Functions.plus);
		return movieGradient;
	}
	
	public HashMap<Integer[], Double> getResidue(){
		double rating;
		double predictedRating;
		DoubleMatrix1D  Vy;
		DoubleMatrix1D  Ux;
		if(userFeatureProjection == null || movieFeatureProjection == null){
			return ratings;
		}
		HashMap<Integer[], Double> residue = new HashMap<Integer[], Double>();
        //Object[] ratingsKey  = ratings.keySet().toArray();
		for(Integer[] userItemPair: ratings.keySet()){
			Integer userId = userItemPair[0];
			Integer movieId = userItemPair[1];
			Ux = userTraits.get(userId);
			Vy = movieTraits.get(movieId);
			predictedRating = algebra.mult(Ux,Vy);
			rating = ratings.get(userItemPair); 
			//System.out.format("Actual : %f Predicted :%f\n",rating,predicted);
			residue.put(new Integer[]{userId,movieId}, rating - predictedRating);
		}
		return residue;
	}
	
	private DoubleMatrix2D appendColumn(DoubleMatrix2D U, DoubleMatrix2D U1){
		if(U == null) return U1;
		return DoubleFactory2D.dense.appendRows(U, U1);
	}
	
	public DoubleMatrix2D[] rankOneApproximation(HashMap<Integer[], Double> residue)
			throws Exception{
		
		Random random = new Random();
		
		HashMap<Integer, DoubleMatrix1D> userTraits;
		HashMap<Integer, DoubleMatrix1D> movieTraits;
		int totalUserFeatures = mlLoader.getUserFeaturesSize() + mlLoader.getUsersCount();
		int totalMovieFeatures = mlLoader.getMovieFeaturesSize() + mlLoader.getMoviesCount();
		DoubleMatrix2D U1 = DoubleFactory2D.dense.random(1,totalUserFeatures) ;
		DoubleMatrix2D V1 = DoubleFactory2D.dense.random(1,totalMovieFeatures) ;
		
//		for(int i=0;i<mlLoader.getUsersCount();i++) {
//			U1.set(0, i, random.nextGaussian());
//		}
//		for(int i=0;i<mlLoader.getMoviesCount();i++) {
//			V1.set(0, i, random.nextGaussian());
//		}
		
		DoubleMatrix2D userGradient;
		DoubleMatrix2D movieGradient;
		DoubleMatrix2D[] result = new DoubleMatrix2D[2];
		double[] diag = new double[totalUserFeatures + totalMovieFeatures];
		
		int[] iprint = {0,0};
		int[] iflag = {0};
		
		double[] variables = new double[totalUserFeatures + totalMovieFeatures];
		
		double[] derivatives = new double[totalUserFeatures + totalMovieFeatures];
		double error = Double.MAX_VALUE;
		double prev_error;
		
		do{
			prev_error = error;
			userTraits = getUserTraits(U1);
		 	movieTraits = getMovieTraits(V1);
			userGradient = getUserGradient(residue, U1, userTraits, movieTraits);
			movieGradient = getMovieGradient(residue, V1, userTraits, movieTraits);
			error = calculateCost(residue, U1, V1,userTraits, movieTraits);
			ColtUtils.mergeFlattenTo1D(U1, V1, variables);
			ColtUtils.mergeFlattenTo1D(userGradient, movieGradient, derivatives);
			LBFGS.lbfgs(variables.length, 5, variables, error, derivatives,
					false, diag, iprint, 0.0001,
					1e-15, iflag);
			ColtUtils.assignByUnflattening(variables,U1,V1);	
		}while(Math.abs(error-prev_error) > 0.0001 && iflag[0] != 0);
		
		result[0] = U1;
		result[1] = V1;
		return result;
	}
	
	public void runGRF() throws Exception{
		DoubleMatrix2D U1;
		DoubleMatrix2D V1;
		DoubleMatrix2D [] uvPair = new DoubleMatrix2D[2];
		HashMap<Integer[], Double> residue = null;
		double[] testError = new double[2];
		for(int i = 0 ; i < dimension ; i++){
			residue = getResidue();
			double norm = 0 ;
			for(Integer[] userItemPair:residue.keySet()){
				norm += Math.abs(residue.get(userItemPair));
			}
			System.out.println("Residue : "+norm);
			uvPair = rankOneApproximation(residue);
			userFeatureProjection = appendColumn(userFeatureProjection, uvPair[0]);
			movieFeatureProjection = appendColumn(movieFeatureProjection, uvPair[1]);
			userTraits = getUserTraits(userFeatureProjection);
		 	movieTraits = getMovieTraits(movieFeatureProjection);
		}
	 	testError = calculateTestError(userTraits, movieTraits);
	 	System.out.println("Test RMSE : "+testError[0]);
	 	System.out.println("Test MAE : "+testError[1]);
		
	}
}
